// Copyright 2019 Christopher Ho
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#ifndef MPC_CONTROLLER__MPC_CONTROLLER_HPP_
#define MPC_CONTROLLER__MPC_CONTROLLER_HPP_

#include <mpc_controller/visibility_control.hpp>
#include <autoware_auto_planning_msgs/msg/trajectory.hpp>
#include <autoware_auto_vehicle_msgs/msg/vehicle_control_command.hpp>
#include <autoware_auto_vehicle_msgs/msg/vehicle_kinematic_state.hpp>

#include <controller_common/controller_base.hpp>
#include <mpc_controller/config.hpp>

#include <ostream>
#include <string>

namespace motion
{
namespace control
{
namespace mpc_controller
{
/// \brief A wrapper around an autogenerated mpc solver for vehicle dynamics control
class MPC_CONTROLLER_PUBLIC MpcController : public controller_common::ControllerBase
{
public:
  constexpr static std::chrono::nanoseconds solver_time_step{std::chrono::milliseconds{100LL}};
  // Disable other constructors, assignment operators: autogenerated code has annoying static
  // variables
  MpcController(const MpcController &) = delete;
  MpcController(MpcController &&) noexcept = delete;
  MpcController & operator=(const MpcController &) = delete;
  MpcController & operator=(MpcController &&) noexcept = delete;

  using Command = autoware_auto_vehicle_msgs::msg::VehicleControlCommand;
  using State = autoware_auto_vehicle_msgs::msg::VehicleKinematicState;
  using Trajectory = autoware_auto_planning_msgs::msg::Trajectory;

  /// \brief Constructor
  explicit MpcController(const Config & config);
  ~MpcController() override = default;

  /// Getter for config class
  const Config & get_config() const;
  /// Setter for config class
  void set_config(const Config & config);

  /// Get trajectory plan computed as a result of solving the optimization problem.
  /// This method triggers the rather expensive copy from the solved vector to the trajectory
  /// representation
  const Trajectory & get_computed_trajectory() const noexcept;
  /// Get derivatives computed from the control problem
  ControlDerivatives get_computed_control_derivatives() const noexcept;

  /// Debug printing
  void debug_print(std::ostream & out) const;

  /// Get name of algorithm, for debugging or diagnostic purposes
  std::string name() const override;

  /// Get name of algorithm, for debugging or diagnostic purposes
  Index get_compute_iterations() const override;

protected:
  /// Checks trajectory
  bool check_new_trajectory(const Trajectory & trajectory) const override;
  /// Sets reference trajectory
  const Trajectory & handle_new_trajectory(const Trajectory & trajectory) override;
  /// Actually compute the control command
  Command compute_command_impl(const State & state) override;

private:
  /// Run the solver subroutine
  MPC_CONTROLLER_LOCAL void solve();
  /// Roll references forward and update weights appropriately, return true if cold start
  MPC_CONTROLLER_LOCAL bool update_references(Index current_idx);
  /// Set initial conditions for problem
  MPC_CONTROLLER_LOCAL void initial_conditions(const Point & state);
  /// Compute delta to roll state forward or back to match first reference
  MPC_CONTROLLER_LOCAL std::chrono::nanoseconds x0_time_offset(const State & state, Index idx);
  /// Compute interpolated command
  MPC_CONTROLLER_LOCAL Command interpolated_command(std::chrono::nanoseconds x0_time_offset);
  /// Ensure that from x0, reference points evolve smoothly, x0 is assumed to be set
  /// if there's a big change to maintain consistency, return true (force cold start)
  MPC_CONTROLLER_LOCAL bool ensure_reference_consistency(Index horizon);
  /// Apply config parameters to the optimization solver, does NOT set m_config
  MPC_CONTROLLER_LOCAL void apply_config(const Config & cfg);
  /// Apply vehicle parameters
  MPC_CONTROLLER_LOCAL void apply_parameters(const VehicleConfig & cfg) noexcept;
  /// Apply upper/lower bounds to state and control variables
  MPC_CONTROLLER_LOCAL void apply_bounds(const LimitsConfig & cfg) noexcept;
  /// Apply weights to optimization objective
  MPC_CONTROLLER_LOCAL void apply_weights(const OptimizationConfig & cfg);
  /// Apply nominal weights for a specified index range
  MPC_CONTROLLER_LOCAL void apply_nominal_weights(const StateWeight & cfg, Index start, Index end);
  /// Apply terminal weight somewhere in the middle of y, for receding horizon case
  MPC_CONTROLLER_LOCAL void apply_terminal_weights(const Index idx);
  /// Takes points from reference trajectory and applies to refernece
  MPC_CONTROLLER_LOCAL
  void set_reference(const Trajectory & traj, Index y_start, Index traj_start, Index count);
  /// Zero out fields of the nominal weights
  MPC_CONTROLLER_LOCAL void zero_nominal_weights(Index start, Index end);
  /// Zero out terminal weights
  MPC_CONTROLLER_LOCAL void zero_terminal_weights() noexcept;
  /// Set terminal weights
  MPC_CONTROLLER_LOCAL void set_terminal_weights(const StateWeight & cfg) noexcept;
  /// Set terminal reference
  MPC_CONTROLLER_LOCAL void set_terminal_reference(const Point & pt) noexcept;
  /// Roll solution and reference forward
  MPC_CONTROLLER_LOCAL void advance_problem(Index count);
  /// Fills the last N trajectory reference points with points from the reference trajectory
  MPC_CONTROLLER_LOCAL void backfill_reference(Index count);

  Config m_config;
  Trajectory::UniquePtr m_interpolated_trajectory{nullptr};
  mutable Trajectory m_computed_trajectory;
  Index m_last_reference_index;
};  // class MpcController
}  // namespace mpc_controller
}  // namespace control
}  // namespace motion
#endif  // MPC_CONTROLLER__MPC_CONTROLLER_HPP_
